{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "e5a40081",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ground\tground\tground\tground\tground\tground\tground\tground\tground\tground\tground\tground\t\n",
      "ground\tground\tground\tground\tground\tground\tground\tground\tground\tground\tground\tground\t\n",
      "ground\tground\tground\tground\tground\tground\tground\tground\tground\tground\tground\tground\t\n",
      "ground\ttrap\ttrap\ttrap\ttrap\ttrap\ttrap\ttrap\ttrap\ttrap\ttrap\tterminal\t\n",
      "(0, 0) -> (0, 0) -1\t(0, 0) -> (1, 0) -1\t(0, 0) -> (0, 0) -1\t(0, 0) -> (0, 1) -1\t\n",
      "(0, 1) -> (0, 1) -1\t(0, 1) -> (1, 1) -1\t(0, 1) -> (0, 0) -1\t(0, 1) -> (0, 2) -1\t\n",
      "(0, 2) -> (0, 2) -1\t(0, 2) -> (1, 2) -1\t(0, 2) -> (0, 1) -1\t(0, 2) -> (0, 3) -1\t\n",
      "(0, 3) -> (0, 3) -1\t(0, 3) -> (1, 3) -1\t(0, 3) -> (0, 2) -1\t(0, 3) -> (0, 4) -1\t\n",
      "(0, 4) -> (0, 4) -1\t(0, 4) -> (1, 4) -1\t(0, 4) -> (0, 3) -1\t(0, 4) -> (0, 5) -1\t\n",
      "(0, 5) -> (0, 5) -1\t(0, 5) -> (1, 5) -1\t(0, 5) -> (0, 4) -1\t(0, 5) -> (0, 6) -1\t\n",
      "(0, 6) -> (0, 6) -1\t(0, 6) -> (1, 6) -1\t(0, 6) -> (0, 5) -1\t(0, 6) -> (0, 7) -1\t\n",
      "(0, 7) -> (0, 7) -1\t(0, 7) -> (1, 7) -1\t(0, 7) -> (0, 6) -1\t(0, 7) -> (0, 8) -1\t\n",
      "(0, 8) -> (0, 8) -1\t(0, 8) -> (1, 8) -1\t(0, 8) -> (0, 7) -1\t(0, 8) -> (0, 9) -1\t\n",
      "(0, 9) -> (0, 9) -1\t(0, 9) -> (1, 9) -1\t(0, 9) -> (0, 8) -1\t(0, 9) -> (0, 10) -1\t\n",
      "(0, 10) -> (0, 10) -1\t(0, 10) -> (1, 10) -1\t(0, 10) -> (0, 9) -1\t(0, 10) -> (0, 11) -1\t\n",
      "(0, 11) -> (0, 11) -1\t(0, 11) -> (1, 11) -1\t(0, 11) -> (0, 10) -1\t(0, 11) -> (0, 11) -1\t\n",
      "(1, 0) -> (0, 0) -1\t(1, 0) -> (2, 0) -1\t(1, 0) -> (1, 0) -1\t(1, 0) -> (1, 1) -1\t\n",
      "(1, 1) -> (0, 1) -1\t(1, 1) -> (2, 1) -1\t(1, 1) -> (1, 0) -1\t(1, 1) -> (1, 2) -1\t\n",
      "(1, 2) -> (0, 2) -1\t(1, 2) -> (2, 2) -1\t(1, 2) -> (1, 1) -1\t(1, 2) -> (1, 3) -1\t\n",
      "(1, 3) -> (0, 3) -1\t(1, 3) -> (2, 3) -1\t(1, 3) -> (1, 2) -1\t(1, 3) -> (1, 4) -1\t\n",
      "(1, 4) -> (0, 4) -1\t(1, 4) -> (2, 4) -1\t(1, 4) -> (1, 3) -1\t(1, 4) -> (1, 5) -1\t\n",
      "(1, 5) -> (0, 5) -1\t(1, 5) -> (2, 5) -1\t(1, 5) -> (1, 4) -1\t(1, 5) -> (1, 6) -1\t\n",
      "(1, 6) -> (0, 6) -1\t(1, 6) -> (2, 6) -1\t(1, 6) -> (1, 5) -1\t(1, 6) -> (1, 7) -1\t\n",
      "(1, 7) -> (0, 7) -1\t(1, 7) -> (2, 7) -1\t(1, 7) -> (1, 6) -1\t(1, 7) -> (1, 8) -1\t\n",
      "(1, 8) -> (0, 8) -1\t(1, 8) -> (2, 8) -1\t(1, 8) -> (1, 7) -1\t(1, 8) -> (1, 9) -1\t\n",
      "(1, 9) -> (0, 9) -1\t(1, 9) -> (2, 9) -1\t(1, 9) -> (1, 8) -1\t(1, 9) -> (1, 10) -1\t\n",
      "(1, 10) -> (0, 10) -1\t(1, 10) -> (2, 10) -1\t(1, 10) -> (1, 9) -1\t(1, 10) -> (1, 11) -1\t\n",
      "(1, 11) -> (0, 11) -1\t(1, 11) -> (2, 11) -1\t(1, 11) -> (1, 10) -1\t(1, 11) -> (1, 11) -1\t\n",
      "(2, 0) -> (1, 0) -1\t(2, 0) -> (3, 0) -1\t(2, 0) -> (2, 0) -1\t(2, 0) -> (2, 1) -1\t\n",
      "(2, 1) -> (1, 1) -1\t(2, 1) -> (3, 1) -100\t(2, 1) -> (2, 0) -1\t(2, 1) -> (2, 2) -1\t\n",
      "(2, 2) -> (1, 2) -1\t(2, 2) -> (3, 2) -100\t(2, 2) -> (2, 1) -1\t(2, 2) -> (2, 3) -1\t\n",
      "(2, 3) -> (1, 3) -1\t(2, 3) -> (3, 3) -100\t(2, 3) -> (2, 2) -1\t(2, 3) -> (2, 4) -1\t\n",
      "(2, 4) -> (1, 4) -1\t(2, 4) -> (3, 4) -100\t(2, 4) -> (2, 3) -1\t(2, 4) -> (2, 5) -1\t\n",
      "(2, 5) -> (1, 5) -1\t(2, 5) -> (3, 5) -100\t(2, 5) -> (2, 4) -1\t(2, 5) -> (2, 6) -1\t\n",
      "(2, 6) -> (1, 6) -1\t(2, 6) -> (3, 6) -100\t(2, 6) -> (2, 5) -1\t(2, 6) -> (2, 7) -1\t\n",
      "(2, 7) -> (1, 7) -1\t(2, 7) -> (3, 7) -100\t(2, 7) -> (2, 6) -1\t(2, 7) -> (2, 8) -1\t\n",
      "(2, 8) -> (1, 8) -1\t(2, 8) -> (3, 8) -100\t(2, 8) -> (2, 7) -1\t(2, 8) -> (2, 9) -1\t\n",
      "(2, 9) -> (1, 9) -1\t(2, 9) -> (3, 9) -100\t(2, 9) -> (2, 8) -1\t(2, 9) -> (2, 10) -1\t\n",
      "(2, 10) -> (1, 10) -1\t(2, 10) -> (3, 10) -100\t(2, 10) -> (2, 9) -1\t(2, 10) -> (2, 11) -1\t\n",
      "(2, 11) -> (1, 11) -1\t(2, 11) -> (3, 11) 100\t(2, 11) -> (2, 10) -1\t(2, 11) -> (2, 11) -1\t\n",
      "(3, 0) -> (2, 0) -1\t(3, 0) -> (3, 0) -1\t(3, 0) -> (3, 0) -1\t(3, 0) -> (3, 1) -100\t\n",
      "(4, 12, 4) {}\n"
     ]
    }
   ],
   "source": [
    "def get_state(row, col):\n",
    "    if row != 3:\n",
    "        return 'ground'\n",
    "    if row == 3 and  col == 0:\n",
    "        return 'ground'\n",
    "    if row == 3 and col == 11:\n",
    "        return 'terminal'\n",
    "    \n",
    "    return 'trap'\n",
    "\n",
    "ROWS = 4\n",
    "COLS = 12\n",
    "ACTIONS = ['left', 'right', 'up', 'down']\n",
    "Num_actions = len(ACTIONS)\n",
    "\n",
    "for r in range(ROWS):\n",
    "    for c in range(COLS):\n",
    "        state = get_state(r, c)\n",
    "        print(state, end='\\t')\n",
    "    print()\n",
    "\n",
    "\n",
    "def move(row, col, action):\n",
    "    \"\"\"\n",
    "    如果移动到陷阱或者终点，则返回当前位置和0分\n",
    "    否则返回新位置和-1分\n",
    "\n",
    "    \"\"\"\n",
    "    # if get_state(row, col) in ['trap', 'terminal']:\n",
    "    #     return row, col, 0\n",
    "    \n",
    "    if action == 'up':\n",
    "        row -= 1\n",
    "    elif action == 'down':\n",
    "        row += 1\n",
    "    elif action == 'left':\n",
    "        col -= 1\n",
    "    elif action == 'right':\n",
    "        col += 1\n",
    "\n",
    "    # 不允许走到边界\n",
    "    row = max(0, row)\n",
    "    row = min(3, row)\n",
    "    col = max(0, col)\n",
    "    col = min(11, col)\n",
    "\n",
    "    reward = -1\n",
    "\n",
    "    # 奖励规则\n",
    "    if get_state(row, col) == 'terminal':\n",
    "        reward = 100\n",
    "        return row, col, reward\n",
    "    \n",
    "    if get_state(row, col) == 'trap':\n",
    "        reward = -100\n",
    "        return row, col, reward\n",
    "\n",
    "    return row, col, reward\n",
    "\n",
    "for r in range(ROWS):\n",
    "    for c in range(COLS):\n",
    "        state = get_state(r, c)\n",
    "        if state == 'ground':\n",
    "            for a in ['up', 'down', 'left', 'right']:\n",
    "                new_r, new_c, reward = move(r, c, a)\n",
    "                print(f'({r}, {c}) -> ({new_r}, {new_c}) {reward}', end='\\t')\n",
    "            print()\n",
    "\n",
    "# 动作函数和Qlearning算法一样\n",
    "\n",
    "import numpy as np\n",
    "import random\n",
    "\n",
    "def get_action(row, col, episode=0):\n",
    "    # 动态调整探索率，随着训练进行逐渐减少探索\n",
    "    epsilon = max(0.01, 0.1 - episode/10000)  # 从0.1逐渐减少到0.01\n",
    "    \n",
    "    # 有小概率选择随机动作\n",
    "    if random.random() < epsilon:\n",
    "        return random.randint(0, Num_actions - 1)\n",
    "    \n",
    "    # 否则选择分数最高的动作\n",
    "    return Q[row, col].argmax() % Num_actions\n",
    "\n",
    "\n",
    "# 初始化在每一个格子里采取每个动作的分数，初始化都是0，因为没有任何的知识\n",
    "Q = np.zeros([ROWS, COLS, Num_actions])\n",
    "\n",
    "# 保存历史数据，键是（row，col， action）， 值是（next_row,next_col, reward）\n",
    "history = {}\n",
    "\n",
    "print(Q.shape, history)\n",
    "\n",
    "def get_update(row, col, action, reward, next_row, next_col):\n",
    "    # target为下一个格子的最高分数，这里的计算和下一步的动作无关\n",
    "    target = 0.9 * Q[next_row, next_col].max()\n",
    "\n",
    "    # 加上本步的分数\n",
    "    target += reward\n",
    "\n",
    "    # 计算value\n",
    "    value = Q[row, col, action]\n",
    "\n",
    "    # 计算更新值\n",
    "    update = 0.1 * (target - value)\n",
    "\n",
    "    return update\n",
    "\n",
    "def q_planning():\n",
    "    # 随机选择曾经遇到过的状态动作对\n",
    "    row, col, action = random.choice(list(history.keys()))\n",
    "\n",
    "    # 再获取下一个状态和反馈\n",
    "    next_row, next_col, reward = history[(row, col, action)]\n",
    "\n",
    "    # 计算分数\n",
    "    update = get_update(row, col ,action, reward, next_row, next_col)\n",
    "\n",
    "    # 更新分数\n",
    "    Q[row, col, action] += update # 将上次的分数加上这次的分数，相加后得到新的分数\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "2835dd6b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 0: reward_sum = -118, steps = 19\n",
      "Episode 200: reward_sum = -110, steps = 11\n",
      "Episode 400: reward_sum = -105, steps = 6\n",
      "Episode 600: reward_sum = 81, steps = 20\n",
      "Episode 800: reward_sum = 88, steps = 13\n",
      "Episode 1000: reward_sum = 88, steps = 13\n",
      "Episode 1200: reward_sum = 88, steps = 13\n",
      "Episode 1400: reward_sum = 88, steps = 13\n",
      " S → → ↓ → → → → → → → → \n",
      " → → → → → → → → → → → → \n",
      " ↓ ↓ ↓ ↓ ↓ ↓ ↓ ↓ ↓ ↓ ↓ → \n",
      " ← C C C C C C C C C C G \n"
     ]
    }
   ],
   "source": [
    "epoch = 1500 # 增加训练次数\n",
    "\n",
    "def train():\n",
    "    for epo in range(epoch):\n",
    "        # 初始化状态，从地面起点开始\n",
    "        row = 3  \n",
    "        col = 0  \n",
    "\n",
    "        # 计算反馈的和\n",
    "        reward_sum = 0\n",
    "\n",
    "        steps = 0\n",
    "        max_steps = 100\n",
    "\n",
    "        while get_state(row, col) not in ['trap', 'terminal'] and steps < max_steps:\n",
    "            # 选择动作（使用动态探索率）\n",
    "            epsilon = max(0.01, 0.5 - epo/2000)  # 从0.5逐渐减少到0.01\n",
    "            if random.random() < epsilon:\n",
    "                action = random.randint(0, Num_actions - 1)\n",
    "            else:\n",
    "                action = Q[row, col].argmax() % Num_actions\n",
    "            \n",
    "            # 执行动作\n",
    "            next_row, next_col, reward = move(row, col, ACTIONS[action])\n",
    "            reward_sum += reward\n",
    "\n",
    "            # 计算更新值\n",
    "            target = 0.9 * Q[next_row, next_col].max()\n",
    "            target += reward\n",
    "            value = Q[row, col, action]\n",
    "            update = 0.1 * (target - value)\n",
    "\n",
    "            # 更新Q值\n",
    "            Q[row, col, action] += update\n",
    "\n",
    "            # 将经验存储到模型中\n",
    "            history[(row, col, action)] = (next_row, next_col, reward)\n",
    "\n",
    "            # 进行Q-planning（适度频率）\n",
    "            if len(history) > 5 and random.random() < 0.3:  # 30%概率进行规划\n",
    "                for _ in range(3):\n",
    "                    q_planning()\n",
    "\n",
    "            # 更新当前位置\n",
    "            row, col = next_row, next_col\n",
    "            steps += 1\n",
    "            \n",
    "        # 输出训练进度\n",
    "        if epo % 200 == 0:\n",
    "            print(f\"Episode {epo}: reward_sum = {reward_sum}, steps = {steps}\")\n",
    "\n",
    "train()\n",
    "\n",
    "\n",
    "# 训练完成后，从Q表中提取最优策略\n",
    "for r in range(ROWS):\n",
    "    for c in range(COLS):\n",
    "        history[r, c] = Q[r, c].argmax()  # 选择Q值最大的动作作为最优策略\n",
    "def show_policy(history):\n",
    "    \"\"\"\n",
    "    可视化策略\n",
    "    \"\"\"\n",
    "    symbols = ['↑','↓','←','→']\n",
    "\n",
    "    for r in range(ROWS):\n",
    "        line = ' '\n",
    "        for c in range(COLS):\n",
    "            s = get_state(r, c)\n",
    "            if s == 'terminal': line += 'G '\n",
    "            elif s == 'trap': line += 'C '\n",
    "            elif (r,c) == (0,0): line += 'S '\n",
    "            else: line += symbols[history[r,c]] + ' '\n",
    "        print(line)\n",
    "\n",
    "show_policy(history)\n",
    "\n",
    "\n",
    "def show(row, col, action):\n",
    "    # 创建12x12的地图可视化\n",
    "    graph = []\n",
    "    for r in range(ROWS):\n",
    "        for c in range(COLS):\n",
    "            state = get_state(r, c)\n",
    "            if state == 'start':\n",
    "                graph.append('S ')  # 起点\n",
    "            elif state == 'terminal':\n",
    "                graph.append('G ')  # 终点\n",
    "            elif state == 'trap':\n",
    "                graph.append('T ')  # 陷阱\n",
    "            else:\n",
    "                graph.append('. ')  # 普通地面\n",
    "    \n",
    "    # 在当前位置显示动作\n",
    "    symbols = ['↑','↓','←','→']\n",
    "    graph[row * COLS + col] = symbols[action] + ' '\n",
    "    \n",
    "    # 打印地图\n",
    "    graph_str = ''.join(graph)\n",
    "    for i in range(0, ROWS * 2, 2):  # 每行2个字符（符号+空格）\n",
    "        print(graph_str[i*COLS:(i+2)*COLS])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "04a209dc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      ". . . . . . . . . . . . \n",
      ". . . . . . . . . . . . \n",
      ". . . . . . . . . . . → \n",
      ". T T T T T T T T T T G \n",
      "游戏结束！最终状态: terminal\n"
     ]
    }
   ],
   "source": [
    "import IPython.display as display\n",
    "import time\n",
    "import os\n",
    "\n",
    "def test():\n",
    "    # 起点\n",
    "    row, col = 0, 0\n",
    "    \n",
    "    # 最多玩N步\n",
    "    for _ in range(200):\n",
    "        # 获取当前状态，如果状态是终点或者掉陷阱则终止\n",
    "        if get_state(row, col) in ['trap', 'terminal']:\n",
    "            print(f\"游戏结束！最终状态: {get_state(row, col)}\")\n",
    "            break\n",
    "            \n",
    "        # 选择最优动作\n",
    "        action = history[row, col]  # 使用训练好的策略\n",
    "        \n",
    "        display.clear_output(wait=True)\n",
    "        # 添加延时以便观察\n",
    "        time.sleep(0.1)\n",
    "        show(row, col, action)\n",
    "        \n",
    "        # 执行动作\n",
    "        next_row, next_col, reward = move(row, col, ACTIONS[action])\n",
    "        \n",
    "        # 更新位置\n",
    "        row, col = next_row, next_col\n",
    "        \n",
    "        \n",
    "\n",
    "# 运行测试\n",
    "test()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
